#include "../lib/sgd.h"

#include "iostream"
#include "cmath"
#include "vector"
#include "string"

#define pi (atan(1.0)*4.0)

//方向梯度类型
enum gradient_type_e
{
	Dnull,
	Dx,
	Dy
};

//正态函数参数组
struct gaussian_para
{
	double mu_x, mu_y, sigma_x, sigma_y, rho;
};

double gaussian_distribution(double x, double y, const gaussian_para &p, 
	gradient_type_e mode = Dnull)
{
	double part = -0.5*(pow((x-p.mu_x)/p.sigma_x,2) -2*p.rho*(x-p.mu_x)*(y-p.mu_y)/(p.sigma_x*p.sigma_y)
		+ pow((y-p.mu_y)/p.sigma_y,2)) / (1.0-p.rho*p.rho);

	double part2;
	if (mode == Dnull)
	{
		return exp(part)/(2*pi*p.sigma_x*p.sigma_y*sqrt(1-p.rho*p.rho));
	}
	else if (mode == Dx)
	{
		part2 = -1.0*((x-p.mu_x)/(p.sigma_x*p.sigma_x) - p.rho*(y-p.mu_y)/(p.sigma_x*p.sigma_y))/(1.0-p.rho*p.rho);
		return part2*exp(part)/(2*pi*p.sigma_x*p.sigma_y*sqrt(1-p.rho*p.rho));
	}
	else if (mode == Dy)
	{
		part2 = -1.0*((y-p.mu_y)/(p.sigma_y*p.sigma_y) - p.rho*(x-p.mu_x)/(p.sigma_x*p.sigma_y))/(1.0-p.rho*p.rho);
		return part2*exp(part)/(2*pi*p.sigma_x*p.sigma_y*sqrt(1-p.rho*p.rho));
	}
	else
	{
		std::string err_str = "Invalid calculation mode.";
		throw err_str;
	}
}

// 正演目标函数的参数
// 声明我们所使用的二维高斯分布的参数
static std::vector<gaussian_para> para;

// 计算当前解的目标函数值以及目标函数在不同方向的梯度
double evaluate(void *instance, const double *x, double *g, const int n, const int m)
{
	double fx = 0.0;
	g[0] = g[1] = 0.0;

	for (int i = 0; i < para.size(); i++)
	{
		fx += -1.0*gaussian_distribution(x[0], x[1], para[i]);
		g[0] += -1.0*gaussian_distribution(x[0], x[1], para[i], Dx);
		g[1] += -1.0*gaussian_distribution(x[0], x[1], para[i], Dy);
	}

	fx += 15.78257991; // 取5组参数 非凸的情况
	//fx += 12.82906044; // 取2组参数 凸的情况

	g[0] *= 2.0*fx;
	g[1] *= 2.0*fx;
	return fx*fx;
}

int progress(void *instance, sgd_float fx, const sgd_float *x, const sgd_float *g, 
	const sgd_para *param, const int n_size, const int k)
{
	std::clog << "iteration time: " << k << ", fx: " << fx << "\r";
	if (fx < param->epsilon) {std::clog << std::endl;}
	return 0;
}

int main(int argc, char *argv[])
{
	gaussian_para tmp_p;

	tmp_p.mu_x = 0.50; tmp_p.mu_y = 0.40; tmp_p.sigma_x = 0.15; tmp_p.sigma_y = 0.10; tmp_p.rho = 0.2;
	para.push_back(tmp_p);

	tmp_p.mu_x = 0.60; tmp_p.mu_y = 0.70; tmp_p.sigma_x = 0.10; tmp_p.sigma_y = 0.20; tmp_p.rho = 0.0;
	para.push_back(tmp_p);

	tmp_p.mu_x = 0.75; tmp_p.mu_y = 0.20; tmp_p.sigma_x = 0.10; tmp_p.sigma_y = 0.12; tmp_p.rho = 0.5;
	para.push_back(tmp_p);

	tmp_p.mu_x = 0.10; tmp_p.mu_y = 0.40; tmp_p.sigma_x = 0.60; tmp_p.sigma_y = 0.70; tmp_p.rho = 0.1;
	para.push_back(tmp_p);

	tmp_p.mu_x = 0.22; tmp_p.mu_y = 0.66; tmp_p.sigma_x = 0.15; tmp_p.sigma_y = 0.12; tmp_p.rho = -0.2;
	para.push_back(tmp_p);

	sgd_float fx;
	sgd_float x[2] = {0.25, 0.20};

	sgd_para my_para = sgd_default_parameters();
	my_para.iteration = 20000;

	my_para.mu = 0.02;

	my_para.alpha = 0.0008;
	int ret = sgd_solver(evaluate, progress, &fx, &x[0], 2, 1, &my_para, nullptr, SGD_MOMENTUM);
	std::clog << "Momentum return: " << sgd_error_str(ret);
	if (ret > 0) std::clog << "\033[1m\033[32m Successed! \033[0m" << std::endl;
	else std::clog << "\033[1m\033[31m Failed! \033[0m" << std::endl;
	std::clog << "fx = " << fx << " ,initial step = " << my_para.alpha << std::endl;
	std::clog << "model: " << x[0] << " " << x[1] << std::endl << std::endl;

	my_para.alpha = 0.0005;
	x[0] = 0.25; x[1] = 0.20;
	ret = sgd_solver(evaluate, progress, &fx, &x[0], 2, 1, &my_para, nullptr, SGD_NAG);
	std::clog << "NAG return: " << sgd_error_str(ret);
	if (ret > 0) std::clog << "\033[1m\033[32m Successed! \033[0m" << std::endl;
	else std::clog << "\033[1m\033[31m Failed! \033[0m" << std::endl;
	std::clog << "fx = " << fx << " ,initial step = " << my_para.alpha << std::endl;
	std::clog << "model: " << x[0] << " " << x[1] << std::endl << std::endl;

	my_para.alpha = 0.3;
	x[0] = 0.25; x[1] = 0.20;
	ret = sgd_solver(evaluate, progress, &fx, &x[0], 2, 1, &my_para, nullptr, SGD_ADAGRAD);
	std::clog << "Adagrad return: " << sgd_error_str(ret);
	if (ret > 0) std::clog << "\033[1m\033[32m Successed! \033[0m" << std::endl;
	else std::clog << "\033[1m\033[31m Failed! \033[0m" << std::endl;
	std::clog << "fx = " << fx << " ,initial step = " << my_para.alpha << std::endl;
	std::clog << "model: " << x[0] << " " << x[1] << std::endl << std::endl;

	my_para.alpha = 0.05;
	my_para.beta_2 = 0.95;
	x[0] = 0.25; x[1] = 0.20;
	ret = sgd_solver(evaluate, progress, &fx, &x[0], 2, 1, &my_para, nullptr, SGD_RMSPROP);
	std::clog << "RMSProp return: " << sgd_error_str(ret);
	if (ret > 0) std::clog << "\033[1m\033[32m Successed! \033[0m" << std::endl;
	else std::clog << "\033[1m\033[31m Failed! \033[0m" << std::endl;
	std::clog << "fx = " << fx << " ,initial step = " << my_para.alpha << std::endl;
	std::clog << "model: " << x[0] << " " << x[1] << std::endl << std::endl;

	my_para.alpha = 0.5;
	my_para.beta_1 = 0.95;
	x[0] = 0.25; x[1] = 0.20;
	ret = sgd_solver(evaluate, progress, &fx, &x[0], 2, 1, &my_para, nullptr, SGD_ADAM);
	std::clog << "Adam return: " << sgd_error_str(ret);
	if (ret > 0) std::clog << "\033[1m\033[32m Successed! \033[0m" << std::endl;
	else std::clog << "\033[1m\033[31m Failed! \033[0m" << std::endl;
	std::clog << "fx = " << fx << " ,initial step = " << my_para.alpha << std::endl;
	std::clog << "model: " << x[0] << " " << x[1] << std::endl << std::endl;

	my_para.alpha = 0.053;
	my_para.beta_1 = 0.9;
	my_para.beta_2 = 0.95;
	x[0] = 0.25; x[1] = 0.20;
	ret = sgd_solver(evaluate, progress, &fx, &x[0], 2, 1, &my_para, nullptr, SGD_NADAM);
	std::clog << "Nadam return: " << sgd_error_str(ret);
	if (ret > 0) std::clog << "\033[1m\033[32m Successed! \033[0m" << std::endl;
	else std::clog << "\033[1m\033[31m Failed! \033[0m" << std::endl;
	std::clog << "fx = " << fx << " ,initial step = " << my_para.alpha << std::endl;
	std::clog << "model: " << x[0] << " " << x[1] << std::endl << std::endl;

	my_para.alpha = 0.5;
	my_para.beta_1 = 0.9;
	my_para.beta_2 = 0.999;
	x[0] = 0.25; x[1] = 0.20;
	ret = sgd_solver(evaluate, progress, &fx, &x[0], 2, 1, &my_para, nullptr, SGD_ADAMAX);
	std::clog << "AdaMax return: " << sgd_error_str(ret);
	if (ret > 0) std::clog << "\033[1m\033[32m Successed! \033[0m" << std::endl;
	else std::clog << "\033[1m\033[31m Failed! \033[0m" << std::endl;
	std::clog << "fx = " << fx << " ,initial step = " << my_para.alpha << std::endl;
	std::clog << "model: " << x[0] << " " << x[1] << std::endl << std::endl;

	my_para.alpha = 0.5;
	my_para.beta_1 = 0.9;
	my_para.beta_2 = 0.999;
	x[0] = 0.25; x[1] = 0.20;
	ret = sgd_solver(evaluate, progress, &fx, &x[0], 2, 1, &my_para, nullptr, SGD_ADABELIEF);
	std::clog << "AdaBelief return: " << sgd_error_str(ret);
	if (ret > 0) std::clog << "\033[1m\033[32m Successed! \033[0m" << std::endl;
	else std::clog << "\033[1m\033[31m Failed! \033[0m" << std::endl;
	std::clog << "fx = " << fx << " ,initial step = " << my_para.alpha << std::endl;
	std::clog << "model: " << x[0] << " " << x[1] << std::endl;
	return 0;
}